# !/usr/bin/env python
# -*- coding=utf-8 -*-

import time
import numpy as np
import pandas as pd
import logging
import matplotlib.pyplot as plt
import warnings

from OrderSpliter import DynamicProgrammingSolution, EachTick, LOG

warnings.filterwarnings("ignore")
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
pd.options.display.max_rows = 20


DATA_PATH = r"C:\Users\新田草\Desktop\正仁量化\data\CY1805.csv"
IC2009_DATA_PATH = r"C:\Users\新田草\Desktop\正仁量化\data\IC2009.csv"
RESULT_PATH = r"C:\Users\新田草\Desktop\正仁量化\data\splited_order.csv"
RESULT_PATH_DP = r"C:\Users\新田草\Desktop\正仁量化\data\splited_order_dp.csv"


contract_multiplier = 5


def get_data(file_path, func=open) -> pd.DataFrame:
    """
    从文件中 IO 获取数据
    :param file_path:
    :param func:
    :return:
    """
    LOG.info("正在加载数据...")
    with func(file_path) as f:
        data = pd.read_csv(f)
    LOG.info("数据加载完成！")
    return data


class ColIndex(object):
    """
    列索引容器
    """
    def __init__(self):
        self.Stock = None
        self.Date_ind = None
        self.Time_ind = None
        self.Timestamp_ind = None
        self.LastPx_ind = None
        self.TickVolumeTrade_ind = None
        self.TickValueTrade_ind = None
        self.AskPx1_ind = None
        self.AskVol1_ind = None
        self.BidPx1_ind = None
        self.BidVol1_ind = None


def get_col_index(columns: list):
    col_index = ColIndex()
    for ind, each in enumerate(columns):
        if each == "Stock":
            col_index.Stock = ind
        if each == "Date":
            col_index.Date_ind = ind
        if each == "Time":
            col_index.Time_ind = ind
        if each == "Timestamp":
            col_index.Timestamp_ind = ind
        if each == "LastPx":
            col_index.LastPx_ind = ind
        if each == "TickVolumeTrade":
            col_index.TickVolumeTrade_ind = ind
        if each == "TickValueTrade":
            col_index.TickValueTrade_ind = ind
        if each == "AskPx1":
            col_index.AskPx1_ind = ind
        if each == "AskVol1":
            col_index.AskVol1_ind = ind
        if each == "BidPx1":
            col_index.BidPx1_ind = ind
        if each == "BidVol1":
            col_index.BidVol1_ind = ind
    return col_index


def main():

    # 加载全部数据
    data = get_data(IC2009_DATA_PATH)

    data_cols = data.columns
    data_col_index = get_col_index(data_cols)

    # 测试数据：速度起见, 可先取部分 data 进行测试
    # test_data = data.iloc[:5000, :]
    # times = pd.to_datetime(test_data["Timestamp"], unit="s")
    # times_index = pd.Index(times)
    # times_index_shanghai = times_index.tz_localize('UTC').tz_convert('Asia/Shanghai')
    # test_data = test_data.set_index([times_index_shanghai])
    # test_data_values = test_data.values

    # 统计交易分布情况

    trades = data.TickVolumeTrade.value_counts().index
    counts = data.TickVolumeTrade.value_counts().values
    trades_counts = dict()
    for trade, count in zip(trades, counts):
        trades_counts[trade] = count
    LOG.info("=" * 100)
    LOG.info("level2 数据中每个 Tick 交易量分布情况(交易量, 出现次数)>>")
    for each in trades_counts.items():
        logging.info(each)
    LOG.info("=" * 100)

    # 转换timestamp格式，并调整时区
    times = pd.to_datetime(data["Timestamp"], unit="s")
    times_index = pd.Index(times)
    times_index_shanghai = times_index.tz_localize('UTC').tz_convert('Asia/Shanghai')
    data["Timestamp_time"] = times_index_shanghai

    # 记录前一个 Tick 数据
    pre_row = None
    logging.info("共计需要处理 %s 个 Tick" % len(data))

    output_results = list()
    for ind, row in enumerate(data.values):
        # 跳过第一条数据
        if pre_row is None:
            pre_row = row
            logging.info("新的一个交易日，跳过第一条数据.")
            continue
        # 空白行，进入下一个交易日
        if np.isnan(row[data_col_index.TickVolumeTrade_ind]):
            logging.info("空白行,进入下一个交易日数据.")
            pre_row = None
            continue
        # 特殊情况处理，当TradeValue != 0，而 TradeVolume = 0
        if row[data_col_index.TickVolumeTrade_ind] != 0 and row[data_col_index.TickValueTrade_ind] == 0:
            row[data_col_index.TickVolumeTrade_ind] = 1
            row[data_col_index.TickValueTrade_ind] = row[data_col_index.LastPx_ind] * 1 * contract_multiplier

        # 如果该tick发生交易, VolumeTrade != 0
        if row[data_col_index.TickVolumeTrade_ind] != 0:
            #  逐tick拆分
            # 方法一：传统思路
            # each_tick = EachTick(row, pre_row, data_col_index)
            # each_tick.split_order()
            # each_tick.output()
            # output_results.extend(each_tick.result) # [{} {} {}]

            # 方法二：动态规划思路求解
            dp_solution = DynamicProgrammingSolution(row, pre_row, data_col_index)
            dp_solution.calculate_direction()
            dp_solution.split_trades()
            dp_solution.output()
            output_results.extend(dp_solution.result)

        if (ind + 1) % 1000 == 0:
            LOG.info("已处理 %s 个Tick" % (ind+1))
        pre_row = row

    # 将最终订单拆分结果整理为 DataFrame格式输出
    final_results = np.array(output_results)
    final_results = np.column_stack((["CY1805"]*len(final_results), final_results))
    results_df = pd.DataFrame(final_results)
    results_df.columns = ["Stock", "Date", "Time", "Timestamp",
                          "Tradeprice", "Tradevolume", "Tradevalue",
                          "Direction"]
    times = pd.to_datetime(results_df["Timestamp"], unit="s")
    times_index = pd.Index(times)
    times_index_shanghai = times_index.tz_localize('UTC').tz_convert('Asia/Shanghai')
    results_df["Timestamp_time"] = times_index_shanghai
    results_df.sort_values(by="Timestamp_time")
    results_df.to_csv(RESULT_PATH_DP)


if __name__ == "__main__":
    start = time.time()
    LOG.info("start_time = %s" % start)
    main()
    print("all data, total_time: ", time.time() - start)
    LOG.info("all data, total_time:  = %s" % (time.time() - start))


